import argparse
import subprocess
import sys

def generate_flamegraph(stack_file, graph_file):
    """Generate a flame graph using the given stack file."""
    flamegraph_cmd = [
        './FlameGraph/flamegraph.pl', 
        '--title="CPU flame graph"', 
        stack_file
    ]
    try:
        with open(graph_file, 'w') as out_file:
            subprocess.run(flamegraph_cmd, stdout=out_file, check=True)
    except subprocess.CalledProcessError as e:
        print(f"Error generating flame graph: {e}")
        sys.exit(1)

def main():
    parser = argparse.ArgumentParser(description='Generate CPU flame graphs.')
    parser.add_argument('duration', type=int, help='Duration of the trace in seconds.')
    parser.add_argument('-p', '--pid', type=str, help='Profile process with one or more comma separated PIDs only.')
    parser.add_argument('-L', '--tid', type=str, help='Profile thread with one or more comma separated TIDs only.')
    parser.add_argument('-U', '--user-stacks-only', action='store_true', help='Show stacks from user space only.')
    parser.add_argument('-K', '--kernel-stacks-only', action='store_true', help='Show stacks from kernel space only.')
    parser.add_argument('-F', '--frequency', type=int, help='Sample frequency, Hertz.')
    parser.add_argument('-c', '--count', type=int, help='Sample period, number of events.')
    parser.add_argument('-d', '--delimited', action='store_true', help='Insert delimiter between kernel/user stacks.')
    parser.add_argument('-a', '--annotations', action='store_true', help='Add _[k] annotations to kernel frames.')
    parser.add_argument('-I', '--include-idle', action='store_true', help='Include CPU idle stacks.')
    parser.add_argument('-f', '--folded', action='store_true', help='Output folded format.')
    parser.add_argument('--hash-storage-size', type=int, help='Number of hash keys that can be stored.')
    parser.add_argument('--stack-storage-size', type=int, help='Number of unique stack traces that can be stored.')
    parser.add_argument('-C', '--cpu', type=int, help='CPU number to run profile on.')
    parser.add_argument('--cgroupmap', type=str, help='Trace cgroups in this BPF map only.')
    parser.add_argument('--mntnsmap', type=str, help='Trace mount namespaces in this BPF map only.')
    parser.add_argument('-o', '--output', type=str, required=True, help='Output file for stack traces.')
    parser.add_argument('-g', '--graph', type=str, required=True, help='Output file for flame graph.')

    args = parser.parse_args()

    # Build the profile.py command
    profile_cmd = [
        'python', 'profile.py'
    ]
    
    if args.pid:
        profile_cmd += ['-p', args.pid]
    if args.tid:
        profile_cmd += ['-L', args.tid]
    if args.user_stacks_only:
        profile_cmd.append('-U')
    if args.kernel_stacks_only:
        profile_cmd.append('-K')
    if args.frequency:
        profile_cmd += ['-F', str(args.frequency)]
    if args.count:
        profile_cmd += ['-c', str(args.count)]
    if args.delimited:
        profile_cmd.append('-d')
    if args.annotations:
        profile_cmd.append('-a')
    if args.include_idle:
        profile_cmd.append('-I')
    if args.folded:
        profile_cmd.append('-f')
    if args.hash_storage_size:
        profile_cmd += ['--hash-storage-size', str(args.hash_storage_size)]
    if args.stack_storage_size:
        profile_cmd += ['--stack-storage-size', str(args.stack_storage_size)]
    if args.cpu:
        profile_cmd += ['-C', str(args.cpu)]
    if args.cgroupmap:
        profile_cmd += ['--cgroupmap', args.cgroupmap]
    if args.mntnsmap:
        profile_cmd += ['--mntnsmap', args.mntnsmap]

    profile_cmd += [str(args.duration)]

    # Print the profile.py command for debugging
    print(f"Running command: {' '.join(profile_cmd)}")
    
    # Run profile.py and redirect output to specified file
    with open(args.output, 'w') as output_file:
        try:
            subprocess.run(profile_cmd, stdout=output_file, check=True)
        except subprocess.CalledProcessError as e:
            print(f"Error running profile.py: {e}")
            sys.exit(1)
    
    # Generate flame graph from the stack traces
    generate_flamegraph(args.output, args.graph)

if __name__ == '__main__':
    main()
